# -*- encoding: utf-8 -*-
"""
@File    : eee.py
@Author  : lilong
@Time    : 2023/2/24 5:25 下午
"""

import numpy as np

import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Layer, LSTM, Lambda, Embedding
from tensorflow.keras import backend as K
from tensorflow.keras.models import Model

# 模型的输入是词ID矩阵，形状为[batch_size, seq_len]，
# 必须明确的是：0作为padding的ID，而1作为UNK的ID。也就是0必定在句子的结尾。
# 构造数据时，必须注意0的位置，这里的是2句话，每句话的长度是4，长度不够的进行padding。
x_sample = np.array(
            [[5, 3, 0, 0],
             [2, 2, 8, 0]]
           )

# 这样生成的mask矩阵大小是[batch_size, seq_len, 1]，然后词ID矩阵经过Embedding层后的大小为[batch_size, seq_len, word_size]，
# 这样一来就可以用mask矩阵对输出结果就行处理了。这种写法只是我的习惯，并非就是唯一的标准。
x_mask = Lambda(lambda x: K.cast(K.greater(K.expand_dims(x_sample, 2), 0), 'float32'))(x_sample)
print('x_mask:', x_mask)
# x_mask: tf.Tensor(
# [[[1.]
#   [1.]
#   [0.]
#   [0.]]
#
#  [[1.]
#   [1.]
#   [1.]
#   [0.]]], shape=(2, 4, 1), dtype=float32)

# input_dim：受到输入的维度的限制，output_dim的大小受到后续的维度限制
# x_sample_t = Embedding(input_dim=10, output_dim=5)(x_sample)
# print('embedding:', x_sample_t)
# embedding: tf.Tensor(
# [[[ 0.04486466  0.00936719  0.03012386  0.04542464  0.03457835]
#   [-0.00690521  0.04227466 -0.02510612 -0.03621238 -0.03407208]
#   [ 0.04992676  0.0475871  -0.01127332  0.0038435   0.03936446]
#   [ 0.04992676  0.0475871  -0.01127332  0.0038435   0.03936446]]
#  [[-0.01768273  0.0431792   0.02256468  0.00848768 -0.02654272]
#   [-0.01768273  0.0431792   0.02256468  0.00848768 -0.02654272]
#   [ 0.03059833 -0.0157055   0.04822886  0.00583827  0.022221  ]
#   [ 0.04992676  0.0475871  -0.01127332  0.0038435   0.03936446]]], shape=(2, 4, 5), dtype=float32)


# 模拟embedding：shape=(2, 4, 5)，batch为2，每句话4个token，且每个token维度是5
x_sample_t = np.array(
    [
        [[1, 0, 2, 2, 4],
         [2, 2, 0, 3, 3],
         [1, 9, 2, 2, 2],
         [1, 0, 8, 2, 2]],

        [[1, 4, 2, 2, 2],
         [5, 2, 0, 3, 3],
         [1, 9, 0, 2, 2],
         [1, 2, 8, 3, 2]]
    ]
)

# 反转，以行为单位
seq_lengths = tf.round(tf.reduce_sum(x_mask, axis=1)[:, 0])
print('seq_len:', seq_lengths)
# seq_len: tf.Tensor([2. 3.], shape=(2,), dtype=float32)
seq_lengths = tf.cast(seq_lengths, 'int32')
# seq_len: tf.Tensor([2 3], shape=(2,), dtype=int32)
print('seq_len:', seq_lengths)

# 此处的seq_len为每条mask内部为1的值的数目和
out = tf.reverse_sequence(x_sample_t, seq_lengths, seq_axis=1, batch_axis=0)
print(out)
# 序列的翻转，但对于MASK的token，不进行翻转
# tf.Tensor(
# [[[2 2 0 3 3]
#   [1 0 2 2 4]
#   [1 9 2 2 2]
#   [1 0 8 2 2]],
#
#  [[1 9 0 2 2]
#   [5 2 0 3 3]
#   [1 4 2 2 2]
#   [1 2 8 3 2]]], shape=(2, 4, 5), dtype=int64)


# # 例子2：翻转的例子
# a = tf.constant(
#     [[1, 2, 3],
#      [4, 5, 6],
#      [7, 8, 9]]
# )
# l = tf.constant([2], tf.int64)  # 每一次翻转长度分别为1,3,2.由于a是（3,3）维的，所以l中数值最大只能是3
# x = tf.reverse_sequence(a, l, seq_axis=0, batch_axis=1)  # 以列为单位进行翻转，翻转的是每一行的元素
# y = tf.reverse_sequence(a, l, seq_axis=1, batch_axis=0)  # 以行为单位进行翻转，翻转的是每一列的元素
# with tf.compat.v1.Session() as sess:
#     print(x.numpy())
#     print(y.numpy())
# # 每一列上的元素种类没有发生变化，但是从每一行来看，行的顺序分别翻转了前1，前2，前3个元素
# [[1 8 6]
#  [4 5 3]
#  [7 2 9]]
# # 每一行上的元素种类没有发生变化，但是从每一列来看，列的顺序分别翻转了前1，前2，前3个元素
# [[1 2 3]
#  [6 5 4]
#  [8 7 9]]


# # 模型的输入是词ID矩阵，形状为[batch_size, seq_len]，
# # 必须明确，其中0作为padding的ID，而1作为UNK的ID。也就是0必定在句子的结尾
# # 这里的是2句话，每句话的长度是4
# x_sample = np.array(
#             [[5, 3, 0, 0],
#              [2, 2, 8, 0]]
#            )
# x_mask = Lambda(lambda x: K.cast(K.greater(K.expand_dims(x_sample, 2), 0), 'float32'))(x_sample)
# print('x_mask:', x_mask)
# # x_mask: tf.Tensor(
# # [[[1.]
# #   [1.]
# #   [0.]
# #   [0.]],
# #
# #  [[1.]
# #   [1.]
# #   [1.]
# #   [0.]]], shape=(2, 4, 1), dtype=float32)
